#include <atomic>
#include <cassert>
#include "coroutine.hpp"

namespace jizhe {

/// 全局静态变量，用于生成协程id
static std::atomic<uint64_t> s_co_id{1};
/// 全局静态变量，用于统计当前的协程数
static std::atomic<uint64_t> s_co_count{0};

/// 线程局部变量，当前线程正在运行的协程
static thread_local Coroutine *t_running_co = nullptr;
/// 线程局部变量，当前线程的主协程，切换到这个协程，就相当于切换到了主线程中运行，智能指针形式
static thread_local Coroutine::ptr t_schedule_co = nullptr;

#define default_stack_size 128 * 1024;

uint64_t Coroutine::GetRunningCoID() {
    if (t_running_co) {
        return t_running_co->getId();
    }
    return 0;
}

Coroutine::Coroutine() {
    SetThis(this);
    m_state = RUNNING;

    getcontext(&m_ctx);

    ++s_co_count;
    m_id = s_co_id++; // 协程id从1开始，用完加1
    LOG_CO("Create schedule Coroutine id = " << m_id);
}

void Coroutine::SetThis(Coroutine *f) {
    t_running_co = f;
}

/**
 * 获取当前协程，同时充当初始化当前线程主协程的作用，这个函数在使用协程之前要调用一下
 */
Coroutine::ptr Coroutine::GetThis() {
    if (t_running_co) {
        return t_running_co->shared_from_this();
    }
    Coroutine::ptr main_fiber(new Coroutine);
    assert(t_running_co == main_fiber.get());
    t_schedule_co = main_fiber;
    return t_running_co->shared_from_this();
}

/**
 * 带参数的构造函数用于创建其他协程，需要分配栈
 */
Coroutine::Coroutine(std::function<void()> cb, size_t stacksize)
        : m_id(s_co_id++)
        , m_cb(cb) {
    ++s_co_count;
    m_stacksize = stacksize ? stacksize : default_stack_size;
    m_stack     = malloc(m_stacksize);

    getcontext(&m_ctx);

    m_ctx.uc_link          = nullptr;
    m_ctx.uc_stack.ss_sp   = m_stack;
    m_ctx.uc_stack.ss_size = m_stacksize;

    makecontext(&m_ctx, &Coroutine::MainFunc, 0);

    LOG_CO("Create task Coroutine id = " << m_id);
}

/**
 * 线程的主协程析构时需要特殊处理，因为主协程没有分配栈和cb
 */
Coroutine::~Coroutine() {
    --s_co_count;
    if (m_stack) {
        // 有栈，说明是子协程，需要确保子协程一定是结束状态
        assert(m_state == TERM);
        free(m_stack);
        LOG_CO("Delete task Coroutine id = " << m_id);
    } else {
        // 没有栈，说明是线程的主协程
        assert(!m_cb);              // 主协程没有cb
        assert(m_state == RUNNING); // 主协程一定是执行状态

        Coroutine *cur = t_running_co; // 当前协程就是自己
        if (cur == this) {
            SetThis(nullptr);
        }
        LOG_CO("Delete schedule Coroutine id = " << m_id);
    }
}

/**
 * 这里为了简化状态管理，强制只有TERM状态的协程才可以重置，但其实刚创建好但没执行过的协程也应该允许重置的
 */
void Coroutine::reset(std::function<void()> cb) {
    assert(m_stack);
    assert(m_state == TERM);
    m_cb = cb;
    getcontext(&m_ctx);

    m_ctx.uc_link          = nullptr;
    m_ctx.uc_stack.ss_sp   = m_stack;
    m_ctx.uc_stack.ss_size = m_stacksize;

    makecontext(&m_ctx, &Coroutine::MainFunc, 0);
    m_state = READY;
}

void Coroutine::resume() {
    assert(m_state != TERM && m_state != RUNNING);
    SetThis(this);
    m_state = RUNNING;

    // LOG_CO("resume");
    swapcontext(&(t_schedule_co->m_ctx), &m_ctx);
}

void Coroutine::yield() {
    /// 协程运行完之后会自动yield一次，用于回到主协程，此时状态已为结束状态
    assert(m_state == RUNNING || m_state == TERM);
    SetThis(t_schedule_co.get());
    if (m_state != TERM) {
        m_state = READY;
    }

    // LOG_CO(" Switch to schedule-coroutine id=" << t_schedule_co->m_id);
    swapcontext(&m_ctx, &(t_schedule_co->m_ctx));
}

void Coroutine::MainFunc() {
    Coroutine::ptr cur = GetThis(); // GetThis()的shared_from_this()方法让引用计数加1
    assert(cur);

    cur->m_cb();
    cur->m_cb    = nullptr;
    cur->m_state = TERM;

    // 手动让t_running_co的引用计数减1
    // 协程运行完毕自动调用析构函数
    auto raw_ptr = cur.get(); 
    cur.reset();
    raw_ptr->yield();

}

} // namespace jizhe